# -*- coding:utf-8  -*-
# @Author: ChenYangMing
# @Time: 2024-08-30

from sklearn.ensemble import RandomForestClassifier

from Common.config_utils import Config


class RandomForestModel:
    """
    随机森林模型，用于分类任务。
    :param random_state: int 或 None，随机数种子，用于结果的可复现性。
    """
    def __init__(
            self,
            random_state=None,
    ):
        """
        初始化随机森林模型。
        :param random_state: 随机数种子，用于结果的可复现性。
        """
        n_estimators = Config.N_ESTIMATORS  # 决策树的数量
        max_depth = Config.MAX_DEPTH  # 决策树的最大层数
        self.model = RandomForestClassifier(n_estimators=n_estimators, max_depth=max_depth,
                                            random_state=random_state)

    def train(
            self,
            x_train,
            y_train,
    ):
        """
        训练随机森林模型。
        :param x_train: numpy.ndarray 或 pandas.DataFrame，训练数据的特征。
        :param y_train: numpy.ndarray，训练数据的目标标签。
        """
        self.model.fit(x_train, y_train)

    def predict(
            self,
            x_test,
    ):
        """
        使用训练好的模型进行预测。
        :param x_test: numpy.ndarray 或 pandas.DataFrame，测试数据的特征。
        :return: numpy.ndarray，模型预测的标签。
        """
        return self.model.predict(x_test)


# 调用示例
# rf_model = RandomForestModel()  # 实例化模型
# rf_model.train(x_train, y_train)  # 训练模型
